#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging
import json
from tencentcloud.common import credential
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
from tencentcloud.vpc.v20170312 import vpc_client, models

logger = logging.getLogger(__name__)

class TencentCloudSecurityGroupManager:
    def __init__(self, secret_id: str, secret_key: str, region: str):
        """
        初始化腾讯云安全组管理器
        
        Args:
            secret_id: 腾讯云密钥ID
            secret_key: 腾讯云密钥
            region: 区域
        """
        self.cred = credential.Credential(secret_id, secret_key)
        self.client = vpc_client.VpcClient(self.cred, region)
        self.region = region
        
    def get_existing_rules(self, security_group_id: str, description: str) -> list:
        """
        获取指定安全组中由本程序添加的规则
        
        Args:
            security_group_id: 安全组ID
            description: 规则描述标识
            
        Returns:
            匹配的规则列表
        """
        try:
            req = models.DescribeSecurityGroupPoliciesRequest()
            params = {
                "SecurityGroupId": security_group_id
            }
            req.from_json_string(json.dumps(params))
            
            resp = self.client.DescribeSecurityGroupPolicies(req)
            
            matched_rules = []
            # 检查入口规则
            if resp.SecurityGroupPolicySet.Ingress:
                for policy in resp.SecurityGroupPolicySet.Ingress:
                    if hasattr(policy, 'PolicyDescription') and policy.PolicyDescription == description:
                        matched_rules.append(policy)
                        
            return matched_rules
        except TencentCloudSDKException as e:
            logger.error(f"获取腾讯云安全组规则失败: {e}")
            return []
        except Exception as e:
            logger.error(f"获取腾讯云安全组规则失败: {e}")
            return []
            
    def remove_rules(self, security_group_id: str, rules: list):
        """
        删除指定的安全组规则
        
        Args:
            security_group_id: 安全组ID
            rules: 要删除的规则列表
        """
        if not rules:
            return
            
        try:
            req = models.DeleteSecurityGroupPoliciesRequest()
            params = {
                "SecurityGroupId": security_group_id,
                "SecurityGroupPolicySet": {
                    "Ingress": []
                }
            }
            
            for rule in rules:
                policy = {
                    "Protocol": rule.Protocol,
                    "Port": rule.Port,
                    "CidrBlock": rule.CidrBlock,
                    "Action": rule.Action,
                    "PolicyDescription": rule.PolicyDescription
                }
                params["SecurityGroupPolicySet"]["Ingress"].append(policy)
                
            req.from_json_string(json.dumps(params))
            self.client.DeleteSecurityGroupPolicies(req)
            
            for rule in rules:
                logger.info(f"成功删除腾讯云安全组规则: {rule.CidrBlock}")
        except TencentCloudSDKException as e:
            logger.error(f"删除腾讯云安全组规则失败: {e}")
        except Exception as e:
            logger.error(f"删除腾讯云安全组规则失败: {e}")
                
    def add_rules(self, security_group_id: str, ips: set, port_range: str, description: str):
        """
        添加新的安全组规则
        
        Args:
            security_group_id: 安全组ID
            ips: IP地址集合
            port_range: 端口范围
            description: 规则描述
        """
        if not ips:
            return
            
        try:
            req = models.CreateSecurityGroupPoliciesRequest()
            params = {
                "SecurityGroupId": security_group_id,
                "SecurityGroupPolicySet": {
                    "Ingress": []
                }
            }
            
            # 转换端口范围格式
            # 腾讯云使用不同的端口格式，需要转换
            protocol = "TCP"
            port = port_range
            if port_range == "1/65535":
                port = "ALL"
            
            for ip in ips:
                policy = {
                    "Protocol": protocol,
                    "Port": port,
                    "CidrBlock": f"{ip}/32",
                    "Action": "ACCEPT",
                    "PolicyDescription": description
                }
                params["SecurityGroupPolicySet"]["Ingress"].append(policy)
                
            req.from_json_string(json.dumps(params))
            self.client.CreateSecurityGroupPolicies(req)
            
            for ip in ips:
                logger.info(f"成功添加腾讯云安全组规则: {ip}")
        except TencentCloudSDKException as e:
            logger.error(f"添加腾讯云安全组规则失败: {e}")
        except Exception as e:
            logger.error(f"添加腾讯云安全组规则失败: {e}")
                
    def sync_rules(self, security_group_id: str, ips: set, port_range: str = "ALL", description: str = "Auto-synced rule"):
        """
        同步安全组规则
        
        Args:
            security_group_id: 安全组ID
            ips: 当前IP地址集合
            port_range: 端口范围
            description: 规则描述标识
        """
        # 获取现有规则
        existing_rules = self.get_existing_rules(security_group_id, description)
        logger.info(f"找到 {len(existing_rules)} 条现有规则需要删除")
        
        # 删除旧规则
        self.remove_rules(security_group_id, existing_rules)
        
        # 添加新规则
        logger.info(f"准备添加 {len(ips)} 条新规则")
        self.add_rules(security_group_id, ips, port_range, description)